#!/usr/bin/env python
import argparse
import numpy as np
from astropy.io import fits
from scipy.stats import sigmaclip


def get_ptc(img1, img2):
    clip = sigmaclip(img1 - img2, 5, 5)[0]
    return np.median(img1 + img2) / 2, clip.var() / 2


def measure_gain(flist1, flist2, xcent=None, ycent=None, xsize=200, ysize=200):
    with fits.open(flist1[0]) as img1, fits.open(flist2[0]) as img2:
        # determine cutout region
        ny, nx = img1[0].data.shape
        if xcent is None:
            xmin, xmax = 0, nx
        else:
            xmin, xmax = xcent - xsize // 2, xcent + xsize // 2
        if ycent is None:
            ymin, ymax = 0, ny
        else:
            ymin, ymax = ycent - ysize // 2, ycent + ysize // 2

        # bias
        bias, bias_var = get_ptc(img1[0].data[ymin:ymax, xmin:xmax],
                                 img2[0].data[ymin:ymax, xmin:xmax])

    n = len(flist1) - 1
    flux, fvar = np.zeros(n), np.zeros(n)
    for i in range(n):
        with fits.open(flist1[i + 1]) as img1, fits.open(flist2[i + 1]) as img2:
            d, v = get_ptc(img1[0].data[ymin:ymax, xmin:xmax],
                           img2[0].data[ymin:ymax, xmin:xmax])
            flux[i] = d - bias
            fvar[i] = v - bias_var

    gain = flux / fvar
    return '{:.3f}'.format(np.median(gain))


def _run(flist, key='IMGINDEX', **kwargs):
    ids = []
    for f in flist:
        with fits.open(f) as img:
            if key in img[0].header:
                ids.append(img[0].header[key])
            else:
                raise Exception('no valid keyword for grouping')
    flist = np.asarray(flist)
    ids = np.asarray(ids)
    uid = np.unique(ids)
    flist1, flist2 = [], []
    for i in uid:
        index = ids == i
        if index.sum() < 2:
            continue
        flist1.append(flist[index][-2])
        flist2.append(flist[index][-1])
    if len(flist1) > 0:
        res = measure_gain(flist1, flist2, **kwargs)
        print('gain = {}'.format(res))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('names', type=str, nargs='*',
                        help='input image names')
    parser.add_argument('--xcent', '-x', default=None, type=int,
                        help='x position of cutout center. If not set, use the full image')
    parser.add_argument('--ycent', '-y', default=None, type=int,
                        help='y position of cutout center. If not set, use the full image')
    parser.add_argument('--xsize', default=200, type=int,
                        help='x size of cutout, default is 200')
    parser.add_argument('--ysize', default=200, type=int,
                        help='y size of cutout, default is 200')
    parser.add_argument('--keyword', '-k', default='IMGINDEX', type=str,
                        help='fits header keyword for PTC image grouping')
    pars = parser.parse_args()
    _run(pars.names, key=pars.keyword, xcent=pars.xcent, ycent=pars.ycent, xsize=pars.xsize, ysize=pars.ysize)
